#!/usr/bin/env python3
# Copyright (C) 2024 Checkmk GmbH - License: GNU General Public License v2
# This file is part of Checkmk (https://checkmk.com). It is subject to the terms and
# conditions defined in the file COPYING, which is part of this source code package.

# <<<splunk_health>>>
# Overall_state red
# File_monitor_input red
# File_monitor_input Tailreader-0 red
# File_monitor_input Batchreader-0 green
# Data_forwarding red
# Data_forwarding Splunk-2-splunk_forwarding red
# Index_processor green
# Index_processor Index_optimization green
# Index_processor Buckets green
# Index_processor Disk_space green

import enum
from typing import Literal, NewType, TypedDict

import pydantic

from cmk.agent_based.v2 import (
    AgentSection,
    CheckPlugin,
    CheckResult,
    DiscoveryResult,
    Result,
    Service,
    State,
    StringTable,
)


class HealthStatus(enum.StrEnum):
    """Supported splunk health status values."""

    GREEN = "green"
    YELLOW = "yellow"
    RED = "red"


SplunkServiceName = NewType("SplunkServiceName", str)
"""The name of the splunk main service."""

type FeatureHealth = dict[str, HealthStatus]
"""Dictionary that stores the health of each splunk feature related to the main service."""


class HealthItem(pydantic.BaseModel):
    """Object that holds relevant information for splunk health check."""

    model_config = pydantic.ConfigDict()

    name: SplunkServiceName
    """The name of the main splunk service."""
    health: HealthStatus
    """The health status of the main splunk service."""
    features: FeatureHealth = pydantic.Field(default_factory=dict)
    """The health of the splunk service's features."""

    def add_feature_health(self, feature: str, health: HealthStatus) -> None:
        """Add entry into feature health mapping."""
        self.features[feature] = health

    @property
    def sorted_features(self) -> list[tuple[str, HealthStatus]]:
        """Return sorted feature health pairs."""
        return sorted(self.features.items())


def format_service_name(value: str) -> SplunkServiceName:
    """Format service name to replace underscores with spaces."""
    return SplunkServiceName(value.replace("_", " "))


HealthRegistry = dict[SplunkServiceName, HealthItem]
"""A registry for collecting feature health for a given service."""

HealthSection = list[HealthItem]
"""The output generated by the parsing function."""


def parse_splunk_health(string_table: StringTable) -> HealthSection:
    """
    Parse splunk health from agent output.

    Note: items will be skipped if a `KeyError` or `ValueError` are raised. A key error is expected
    if a feature item is parsed without the main service item present, while a value error is raised
    when an invalid value is passed as an attribute to `HealthItem`.
    """
    registry: HealthRegistry = {}

    for state_detail in string_table:
        try:
            match state_detail:
                case (raw_name, health):
                    name = format_service_name(raw_name)
                    registry[name] = HealthItem(name=name, health=HealthStatus(health))
                case (raw_name, feature_name, health):
                    name = format_service_name(raw_name)
                    registry[name].add_feature_health(feature_name, HealthStatus(health))
        except (KeyError, ValueError):
            continue

    return list(registry.values())


def discover_splunk_health(section: HealthSection) -> DiscoveryResult:
    """Discovers splunk health services from parsed agent section."""
    yield Service()


type StateValue = Literal[0, 1, 2, 3]
"""A valid integer code related to the result state."""


class CheckParams(TypedDict):
    """Parameters passed to plugin via ruleset (see defaults)."""

    green: StateValue
    yellow: StateValue
    red: StateValue


def check_splunk_health(params: CheckParams, section: HealthSection) -> CheckResult:
    """Checks the splunk health section returning valid checkmk results."""
    for item in section:
        state = State(params[item.health.value])
        summary = f"{item.name}: {item.health}"
        details = "\n".join(f"{feature}: {health}" for feature, health in item.sorted_features)

        yield Result(state=state, summary=summary, details=details or None)


agent_section_splunk_health = AgentSection(
    name="splunk_health",
    parse_function=parse_splunk_health,
)

check_plugin_splunk_health = CheckPlugin(
    name="splunk_health",
    service_name="Splunk Health",
    check_function=check_splunk_health,
    discovery_function=discover_splunk_health,
    check_ruleset_name="splunk_health",
    check_default_parameters=CheckParams(
        green=State.OK.value,
        yellow=State.WARN.value,
        red=State.CRIT.value,
    ),
)
